from concurrent.futures import (
    ProcessPoolExecutor,
    ThreadPoolExecutor
)
import aiomultiprocess
import functools
import asyncio
from typing import Callable
from enum import Enum, IntEnum
from typing import Optional
import os

from settings.config import settings
from common.log import log

class AsyncProcessPool:
    """多进程算法执行器（单例模式）/后台执行
    """
    
    def __init__(self, max_concurrent_tasks: Optional[int] = None, max_workers: Optional[int] = None) -> None:
        self._max_concurrent_tasks = max_concurrent_tasks or settings.MAX_CONCURRENT_TASKS
        self._semaphore = asyncio.Semaphore(self._max_concurrent_tasks)
        process_workers = max_workers or (os.cpu_count() or 1)
        self._executor = ProcessPoolExecutor(process_workers)            
            
    async def submit(self, task_func: Callable, *args, **kwargs) -> None:
        """提交算法任务（CPU密集型任务）

        Args:
            task_func (Callable): 算法任务
        Returns:
            None
        """
        async with self._semaphore:
            loop = asyncio.get_running_loop()
            task = asyncio.create_task(loop.run_in_executor(self._executor, task_func, *args, **kwargs))
            task.add_done_callback(self._handle_task_result)

    async def shutdown(self, wait: bool = True) -> None:
        await asyncio.to_thread(self._executor.shutdown, wait)
        
    def _handle_task_result(self, task: asyncio.Task) -> None:
        try:
            task.result()   # 获取任务结果
        except Exception as e:
            log.error(f"任务执行异常: {e}")
        
    def __del__(self):
        self._executor.shutdown(wait=True)

class AsyncThreadPool:
    """
    异步线程池工具类，用于执行同步阻塞任务。
    """
    
    def __init__(self, max_workers: Optional[int] = None):
        """
        初始化线程池。

        :param max_workers: 线程池的最大工作线程数，默认为 `os.cpu_count()` 的值。
        """
        self.max_workers = max_workers or (os.cpu_count() or 1)
        self.executor = ThreadPoolExecutor(max_workers=self.max_workers)

    async def run_in_thread(self, func: Callable, timeout: Optional[float] = None, *args, **kwargs):
        """
        在线程池中异步执行一个函数。

        :param func: 要执行的函数。
        :param args: 函数的位置参数。
        :param kwargs: 函数的关键字参数。
        :return: 函数的返回值。
        """
        loop = asyncio.get_running_loop()
        wrapped_func = functools.partial(func, *args, **kwargs)
        try:
            # 添加超时控制
            if timeout is not None:
                return await asyncio.wait_for(loop.run_in_executor(self.executor, wrapped_func), timeout=timeout)
            else:
                return await loop.run_in_executor(self.executor, wrapped_func)
        except asyncio.TimeoutError:
            log.warning(f"任务超时: {func.__name__}")
            raise
        except Exception as e:
            log.error(f"任务执行失败: {e}")
            raise

    async def shutdown(self, wait: bool=True):
        """
        关闭线程池。

        :param wait: 是否等待所有任务完成后再关闭线程池，默认为 `True`。
        """
        await asyncio.to_thread(self.executor.shutdown, wait)